package com.pestdetection;

import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.List;
import java.util.ArrayList;
import java.util.Collections;

/**
 * 病虫害检测器类
 * 用于加载训练好的模型并进行病虫害检测
 * 主要功能：
 * 1. 加载训练好的模型
 * 2. 加载和预处理图片
 * 3. 进行病虫害检测
 * 4. 输出检测结果
 */
public class PestDetector {
    private final MultiLayerNetwork 模型;  // 训练好的模型
    private final NativeImageLoader 图片加载器;  // 用于加载图片
    private final DataNormalization 标准化器;  // 用于标准化图片数据
    private static final int 图片高度 = 224;  // 图片高度
    private static final int 图片宽度 = 224;  // 图片宽度
    private static final int 图片通道数 = 3;  // RGB彩色图片
    
    // 病虫害种类标签映射
    private final Map<Integer, String> 种类标签;  // 索引到名称的映射
    private final Map<String, Integer> 反向种类标签;  // 名称到索引的映射

    /**
     * 构造函数
     * @param 模型路径 模型文件所在路径
     * @throws IOException 如果加载模型失败
     */
    public PestDetector(String 模型路径) throws IOException {
        // 第一步：加载训练好的模型
        this.模型 = MultiLayerNetwork.load(new File(模型路径), true);
        
        // 第二步：创建图片加载器
        this.图片加载器 = new NativeImageLoader(图片高度, 图片宽度, 图片通道数);
        
        // 第三步：创建数据标准化器
        this.标准化器 = new ImagePreProcessingScaler(0, 1);
        
        // 第四步：加载病虫害种类标签
        Map<Integer, String> 标签 = 加载种类标签("dataset/train");
        this.种类标签 = 标签;
        
        // 第五步：创建反向映射（用于名称到索引的查找）
        this.反向种类标签 = new HashMap<>();
        for (Map.Entry<Integer, String> 标签项 : 标签.entrySet()) {
            反向种类标签.put(标签项.getValue(), 标签项.getKey());
        }
    }

    /**
     * 从训练数据目录加载病虫害种类标签
     * @param 数据路径 训练数据所在目录
     * @return 种类标签映射
     */
    private Map<Integer, String> 加载种类标签(String 数据路径) throws IOException {
        Map<Integer, String> 标签 = new HashMap<>();
        File 数据目录 = new File(数据路径);
        
        if (!数据目录.exists() || !数据目录.isDirectory()) {
            throw new IOException("训练数据目录不存在: " + 数据路径);
        }

        // 获取所有子目录（每个子目录代表一种病虫害）
        File[] 种类目录 = 数据目录.listFiles(File::isDirectory);
        if (种类目录 == null || 种类目录.length == 0) {
            throw new IOException("未找到病虫害种类目录");
        }

        // 按字母顺序排序种类名称
        List<String> 种类名称列表 = new ArrayList<>();
        for (File 目录 : 种类目录) {
            种类名称列表.add(目录.getName());
        }
        Collections.sort(种类名称列表);

        // 创建标签映射
        for (int 索引 = 0; 索引 < 种类名称列表.size(); 索引++) {
            标签.put(索引, 种类名称列表.get(索引));
        }

        System.out.println("加载病虫害种类标签：");
        for (Map.Entry<Integer, String> 标签项 : 标签.entrySet()) {
            System.out.println(String.format("- %d: %s", 标签项.getKey(), 标签项.getValue()));
        }

        return 标签;
    }

    /**
     * 检测单张图片
     * @param 图片路径 要检测的图片路径
     * @return 检测结果，包含病虫害类型和置信度
     * @throws IOException 如果处理图片失败
     */
    public DetectionResult detect(String 图片路径) throws IOException {
        // 第一步：检查图片文件是否存在
        File 图片文件 = new File(图片路径);
        if (!图片文件.exists()) {
            throw new IOException("图片文件不存在: " + 图片路径);
        }

        // 第二步：加载并预处理图片
        INDArray 图片数据 = 图片加载器.asMatrix(图片文件);
        标准化器.transform(图片数据);
        
        // 第三步：使用模型进行预测
        INDArray 预测结果 = 模型.output(图片数据);
        
        // 第四步：解析预测结果
        int 预测种类索引 = 预测结果.argMax(1).getInt(0);
        double 置信度 = 预测结果.getDouble(预测种类索引);
        String 预测种类名称 = 种类标签.get(预测种类索引);
        
        // 第五步：将英文标签转换为中文显示
        String 显示名称 = 获取显示名称(预测种类名称);
        
        return new DetectionResult(
            显示名称,
            置信度,
            !预测种类名称.equals("healthy")  // 如果不是healthy，则判定为有病虫害
        );
    }

    /**
     * 获取所有病虫害种类的概率
     * @param 图片路径 要检测的图片路径
     * @return 所有种类的概率
     * @throws IOException 如果处理图片失败
     */
    public Map<String, Double> getAllProbabilities(String 图片路径) throws IOException {
        // 第一步：检查图片文件是否存在
        File 图片文件 = new File(图片路径);
        if (!图片文件.exists()) {
            throw new IOException("图片文件不存在: " + 图片路径);
        }

        // 第二步：加载并预处理图片
        INDArray 图片数据 = 图片加载器.asMatrix(图片文件);
        标准化器.transform(图片数据);
        
        // 第三步：使用模型进行预测
        INDArray 预测结果 = 模型.output(图片数据);
        
        // 第四步：获取所有种类的概率
        Map<String, Double> 概率分布 = new HashMap<>();
        for (int 索引 = 0; 索引 < 预测结果.length(); 索引++) {
            String 种类名称 = 种类标签.get(索引);
            String 显示名称 = 获取显示名称(种类名称);
            概率分布.put(显示名称, 预测结果.getDouble(索引));
        }
        
        return 概率分布;
    }

    /**
     * 将英文标签转换为中文显示
     * @param 标签 英文标签
     * @return 中文显示名称
     */
    private String 获取显示名称(String 标签) {
        switch (标签.toLowerCase()) {
            case "baifenbing":
                return "白粉病";
            case "huanglongbing":
                return "黄龙病";
            case "tanjubing":
                return "炭疽病";
            case "healthy":
                return "健康";
            default:
                return 标签;
        }
    }

    /**
     * 检测结果类
     * 包含病虫害类型、置信度和是否有病虫害的信息
     */
    public static class DetectionResult {
        private final String 病虫害类型;  // 检测到的病虫害类型
        private final double 置信度;      // 检测结果的置信度
        private final boolean 有病虫害;   // 是否检测到病虫害

        public DetectionResult(String 病虫害类型, double 置信度, boolean 有病虫害) {
            this.病虫害类型 = 病虫害类型;
            this.置信度 = 置信度;
            this.有病虫害 = 有病虫害;
        }

        public String getPestType() {
            return 病虫害类型;
        }

        public double getConfidence() {
            return 置信度;
        }

        public boolean hasPest() {
            return 有病虫害;
        }
    }
} 